import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import matplotlib
font = {'family': 'SimHei'}
matplotlib.rc('font', **font)


# 构建模型
def sigmoid(x, w, A):
    return A/(1+np.exp(-w*(x-11)))
def d_x(x, w, A):
    s = sigmoid(x, w, A)
    return w*s*(1-s/A)
def d_w(x, w, A):     # 对w求导
    return (x-11)*sigmoid(x, w, A)*(1-d_A(x, w, A))
def d_A(x, w, A):     # 对A求导
    return sigmoid(x, w, A)/A


def upgrade_arg(N, lam, data):
    alpha1 = 10
    alpha2 = 1e-7
    m = len(data)
    x0 = np.linspace(0, m-1, 30)
    cost = []
    y = np.mat(data).reshape(-1, 1)
    val = 1e9
    for i in range(10000):
        y_pre = sigmoid(x0, lam, N).reshape(-1, 1)
        error = y_pre - y
        cost_val = 1/(2*m)*error.T.dot(error)
        if cost_val[0, 0] < val:
            cost.append(cost_val[0, 0])
            differ1 = sum(d_A(x0, lam, N))/30
            N += alpha1*differ1
        else:
            break
        val = cost_val[0, 0]
    val = 1e8
    for i in range(10000):
        y_pre = sigmoid(x0, lam, N).reshape(-1, 1)
        error = y_pre - y
        cost_val = 1/(2*m)*error.T.dot(error)
        if cost_val[0, 0] < val:
            cost.append(cost_val[0, 0])
            differ2 = sum(d_w(x0, lam, N))/30
            lam += alpha2*differ2
        else:
            break
        val = cost_val[0, 0]
    plt.plot(cost)  # 展示代价值变化
    plt.xlabel('优化次数')
    plt.ylabel('代价值')
    plt.show()
    N = int(N)
    print('确诊总人数=', N, '\n确诊患者人均感染人数=%.2f' % lam)
    return N, lam


def draw(N, lam):
    x0 = np.linspace(0, len(df['确诊累计'])-1, 30)
    x1 = np.linspace(len(df['确诊累计'])-1, 35, 30)
    fig, ax = plt.subplots(1, 1, figsize=(8, 5))
    ax1 = ax.twinx()
    ax.plot(x0, d_x(x0, lam, N), label='拟合线')
    ax.plot(x1, d_x(x1, lam, N), label='预测线')
    ax.scatter(range(30), df[' all'], color='green', label='新增确诊人数')
    ax.set_ylabel('新增人数')
    ax.annotate('新增', xy=(20, 2100), xytext=(25, 4000), color='purple',
                 arrowprops=dict(arrowstyle="->"), fontsize=14)
    ax.legend(loc='upper left')
    ax1.plot(x0, sigmoid(x0, lam, N))
    ax1.plot(x1, sigmoid(x1, lam, N))
    ax1.scatter(range(30), df['确诊累计'], color='red', label='累计确诊人数')
    ax1.set_ylabel('累计人数')
    ax1.annotate('累计', xy=(23, 70000), xytext=(25, 60000), color='purple',
                 arrowprops=dict(arrowstyle="->"), fontsize=14)
    ax1.legend(loc='center left')
    plt.savefig('曲线.jpg')
    plt.show()


# 数据准备
df = pd.read_table('WU_FILE_0.txt', sep='   ', encoding='gbk', engine='python')
df['确诊累计'] = pd.DataFrame([sum(df[' all'].iloc[:i+1]) for i in range(31)])

# 初始参数
N = 40000
lam = 0.2

# 梯度下降优化参数
N, lam = upgrade_arg(N, lam, df['确诊累计'])

# 传入优化好的参数
draw(N, lam)
